#ifndef _BUFFER_H_
#define _BUFFER_H_

#include <iostream>

#include "../glew/glew.h"

#include <cuda_runtime.h>
#include <cuda_gl_interop.h>

#include "../cuCommon/cuTypes.h"
#include "../cuCommon/cuFunction.h"

#define LINEAR_PITCH 16

#define BUFFER_H                       0x01 // Buffer Host memory pointer
#define BUFFER_D                       0x02 // Buffer Device memory pointer
#define BUFFER_A                       0x04 // Buffer Array memory pointer
#define BUFFER_B                       0x08 // Buffer GL object memory pointer
#define BUFFER_V                       0x10 // Buffer Volume memory pointer

#define BUFFER_H_TO_D                  0x0001 // Buffer Host To Device Direction
#define BUFFER_D_TO_H                  0x0002 // Buffer Device To Host Direction

#define BUFFER_H_TO_A                  0x0004 // Buffer Host To Array Direction
#define BUFFER_A_TO_H                  0x0008 // Buffer Array To Host Direction

#define BUFFER_D_TO_A                  0x0010 // Buffer Device to Array Direction
#define BUFFER_A_TO_D                  0x0020 // Buffer Array to Device Direction

#define BUFFER_D_TO_B                  0x0040 // Buffer Device to Buffer Direction
#define BUFFER_B_TO_D                  0x0080 // Buffer Buffer to Device Direction

#define BUFFER_H_TO_B                  0x0100 // Buffer Host to Buffer Direction
#define BUFFER_B_TO_H                  0x0200 // Buffer Buffer to Host Direction

template <typename T>
class TBuffer
{
	T         * m_hptr;
	T         * m_dptr;
	uint32      m_buffer;
	cudaArray * m_array;

	uint32      m_size;
	uint3       m_dim;
	uint3       m_pitch;
	uint3       m_whd;
	uint8       m_flag;

	void Null()
	{
		m_hptr   = NULL;
		m_dptr   = NULL;
		m_array  = NULL;
		m_buffer = 0;

		m_size = 0;
		
		m_dim = ui3();
		m_pitch = ui3();
		m_whd = ui3();

		m_flag = 0;
	}
	
	void Pitch(const uint3 &d)
	{
		m_pitch.x = (d.x % LINEAR_PITCH) ? (LINEAR_PITCH - d.x % LINEAR_PITCH) : 0;
		m_pitch.y = 0;
		m_pitch.z = 0;

		m_whd.x = m_pitch.x + d.x;
		m_whd.y = m_pitch.y + d.y;
		m_whd.z = m_pitch.z + d.z;
	}

	void Create_Host()
	{
		m_hptr = new T[m_size];

#ifdef _DEBUG
		if (m_hptr == NULL) std::cerr << "Host memory allocation failed" << std::endl;
#endif
	}

	void Release_Host()
	{
		delete [] m_hptr;

		m_hptr = NULL;
	}

	void Create_Device()
	{
		cudaMalloc( (void ** ) &m_dptr, m_size * sizeof(T) );

#ifdef _DEBUG
		if (m_dptr == NULL) std::cerr << "Device memory allocation failed" << std::endl;
#endif
	}

	void Release_Device()
	{
		cudaFree(m_dptr);
	}

	void Create_Buffer()
	{
		glGenBuffers(1, &m_buffer);
		glBindBuffer(GL_ARRAY_BUFFER, m_buffer);
		glBufferData(GL_ARRAY_BUFFER, m_size * sizeof(T), NULL, GL_DYNAMIC_DRAW);

#ifdef _DEBUG
		if (m_buffer == 0) std::cerr << "GL Buffer memory allocation failed" << std::endl;
#endif
		cudaGLRegisterBufferObject(m_buffer);
	}

	void Release_Buffer()
	{
		cudaGLUnregisterBufferObject(m_buffer);
		glDeleteBuffers(1, &m_buffer);
	}

	void Create_Array()
	{
		cudaChannelFormatDesc desc = cudaCreateChannelDesc<T>();
		cudaMallocArray( &m_array, &desc, m_whd.x, m_whd.y);

#ifdef _DEBUG
		if (m_array == NULL) std::cerr << "Array memory allocation failed" << std::endl;
#endif
	}

	void Release_Array()
	{
		cudaFreeArray(m_array);
	}

	void Create_Volume()
	{
		cudaExtent extent;

		extent.width = m_whd.x;
		extent.height = m_whd.y;
		extent.depth = m_whd.z;

		cudaChannelFormatDesc desc = cudaCreateChannelDesc<T>();

		cudaMalloc3DArray( &m_array, &desc, extent);

#ifdef _DEBUG
		if (m_array == NULL) std::cerr << "Volume memory allocation failed" << std::endl;
#endif
	}

	void Release_Volume()
	{
		cudaFreeArray(m_array);
	}

public:
	TBuffer()
	{
		Null();
	}

	~TBuffer()
	{
		Release(m_flag);
	}

	TBuffer(const TBuffer & b)
	{
		Clone(*this, b);
	}

  TBuffer & operator = (TBuffer & b)
	{
		Clone(*this, b);

		return *this;
	}

	
	//////////////////////////////////////////////////////////////////////////
	/// Clone function
	//////////////////////////////////////////////////////////////////////////
	static void Clone(TBuffer & dst, TBuffer & src)
	{
		uint32 size = src.size();

		if (size == 0) return;

		uint3 dim = src.dim();

		uint8 createFlag = src.flag();

		dst.Create(createFlag, dim.x, dim.y, dim.z);

		if (createFlag & BUFFER_H) memcpy         (dst.hPtr(), src.hPtr(), size * sizeof(T));
		if (createFlag & BUFFER_D) cudaMemcpyDtoD (dst.dPtr(), src.dPtr(), size * sizeof(T));
		if (createFlag & BUFFER_A) cudaMemcpyAtoA (dst.aPtr(), src.aPtr(), size * sizeof(T));
		if (createFlag & BUFFER_V) cudaMemcpyAtoA (dst.aPtr(), src.aPtr(), size * sizeof(T));
		if (createFlag & BUFFER_B) 
		{
			T * pSrc = NULL, * pDst = NULL;

			cudaGLMapBufferObject((void **) &pSrc, src.bVal());
			cudaGLMapBufferObject((void **) &pDst, dst.bVal());

			cudaMemcpyDtoD (pDst, pSrc, size * sizeof(T));

			cudaGLUnmapBufferObject(src.bVal());
			cudaGLUnmapBufferObject(dst.bVal());
		}
	}

	//////////////////////////////////////////////////////////////////////////
	/// Create function allocates resources
	//////////////////////////////////////////////////////////////////////////
	void Create(uint8 createFlag)
	{
		if (m_size == 0) return;

		if (createFlag & BUFFER_H) Create_Host();
		if (createFlag & BUFFER_D) Create_Device();
		if (createFlag & BUFFER_B) Create_Buffer();
		if (createFlag & BUFFER_A) Create_Array();
		if (createFlag & BUFFER_V) Create_Volume();

		m_flag |= createFlag;
	}

	//////////////////////////////////////////////////////////////////////////
	/// Create function allocates resources
	//////////////////////////////////////////////////////////////////////////
	void Create(uint8 createFlag, uint32 w, uint32 h = 1, uint32 d = 1)
	{
		// first save dimensions
		m_dim = ui3(w, h, d);

		// compute pitch and store results in m_pitch and m_whd
		Pitch(m_dim);

		// m_whd store dimensions + pitch
		m_size = m_whd.x * m_whd.y * m_whd.z;

		if (m_size == 0) return;

		if (createFlag & BUFFER_H) Create_Host();
		if (createFlag & BUFFER_D) Create_Device();
		if (createFlag & BUFFER_B) Create_Buffer();
		if (createFlag & BUFFER_A) Create_Array();
		if (createFlag & BUFFER_V) Create_Volume();

		m_flag |= createFlag;
	}

	void Release(uint8 releaseFlag)
	{
		if (releaseFlag & BUFFER_H) Release_Host();
		if (releaseFlag & BUFFER_D) Release_Device();
		if (releaseFlag & BUFFER_B) Release_Buffer();
		if (releaseFlag & BUFFER_A) Release_Array();
		if (releaseFlag & BUFFER_V) Release_Volume();

		m_flag ^= releaseFlag;
	}

	uint32      size() { return m_size; }
	uint3       whd()  { return m_whd;  }
	uint3       dim()  { return m_dim;  }
	uint8       flag() { return m_flag; }

	T         * hPtr() { return m_hptr;   }
	T         * dPtr() { return m_dptr;   }
	cudaArray * aPtr() { return m_array;  }
	uint32      bVal()  { return m_buffer; }

	void Memcpy(uint32 dirFlag, uint32 start = 0, uint32 count = 0)
	{
		if (count == 0) count = m_size - start;

#ifdef _DEBUG
		if (start + count > m_size) std::cerr << "Invalid start / count arguments" << std::endl;
#endif

		T * ptr = NULL;

		switch (dirFlag)
		{
			case BUFFER_H_TO_D: 
				cudaMemcpyHtoD(m_dptr + start, m_hptr + start, count * sizeof(T));
			break;
			
			case BUFFER_D_TO_H:
				cudaMemcpyDtoH(m_hptr + start, m_dptr + start, count * sizeof(T));
			break;

			case BUFFER_H_TO_B: 
				glBindBuffer(GL_ARRAY_BUFFER, m_buffer);

				ptr = (T *) glMapBuffer(m_buffer, GL_WRITE_ONLY);
				memcpy(ptr + start, m_hptr + start, count * sizeof(T));
				
				glUnmapBuffer(GL_ARRAY_BUFFER);
			break;
			
			case BUFFER_B_TO_H:
				glBindBuffer(GL_ARRAY_BUFFER, m_buffer);

				ptr = (T *) glMapBuffer(m_buffer, GL_READ_ONLY);
				memcpy(m_hptr + start, ptr + start, count * sizeof(T));
				
				glUnmapBuffer(GL_ARRAY_BUFFER);
			break;

			case BUFFER_B_TO_D: 
				//glBindBuffer(GL_ARRAY_BUFFER, m_buffer);

				cudaGLMapBufferObject((void **) &ptr, m_buffer);
				
				cudaMemcpyDtoD(m_dptr + start, ptr + start, count * sizeof(T));

				cudaGLUnmapBufferObject(m_buffer);
				
				//glUnmapBuffer(GL_ARRAY_BUFFER);
			break;
			
			case BUFFER_D_TO_B:
				cudaGLMapBufferObject((void **) &ptr, m_buffer);
				
				cudaMemcpyDtoD(ptr + start, m_dptr + start, count * sizeof(T)); // copy from device to buffer memory

				cudaGLUnmapBufferObject(m_buffer);
			break;

			case BUFFER_H_TO_A: 				
				cudaMemcpyHtoA( m_array, m_hptr, count * sizeof(T));
			break;

			case BUFFER_D_TO_A:
				cudaMemcpyDtoA( m_array, m_dptr, count * sizeof(T));
				break;	

			case BUFFER_A_TO_H:
				cudaMemcpyAtoH( m_hptr, m_array, count * sizeof(T));
			break;

			case BUFFER_A_TO_D:
				cudaMemcpyAtoD( m_dptr, m_array, count * sizeof(T));
				break;

			default: 
				std::cerr << "Invalid direction argument" << std::endl;
			break;
		}
	}

	void Memcpy2D(uint32 dirFlags, uint32 start = 0, uint32 count = 0, uint32 ox = 0, uint32 oy = 0)
	{
		if (count == 0) count = m_size - start;

#ifdef _DEBUG
		if (start + count > m_size) std::cerr << "Invalid start / count arguments" << std::endl;
#endif

		switch (dir)
		{
			//case BUFFER_H_TO_A: 				
			//	cudaMemcpyHtoA( m_array, ox, oy, m_hptr, count * sizeof(T));
			//break;

			//case BUFFER_D_TO_A:
			//	cudaMemcpyDtoA( m_array, ox, oy, m_dptr, count * sizeof(T));
			//	break;	

			//case BUFFER_A_TO_H:
			//	cudaMemcpyAtoH( m_hptr, m_array, ox, oy, count * sizeof(T));
			//break;

			//case BUFFER_A_TO_D:
			//	cudaMemcpyAtoD( m_dptr, m_array, ox, oy, count * sizeof(T));
			//	break;

			default: 
				std::cerr << "Invalid direction argument" << std::endl;
			break;
		}
	}

	void Memset(uint8 bufFlag, T * p, uint32 start = 0, uint32 count = 0)
	{
		if (count == 0) count = m_size - start;

		if (bufFlag & BUFFER_H) memset(m_hptr + start, 0, count * sizeof(T));
		if (bufFlag & BUFFER_D) cudaMemset(m_dptr + start, 0, count * sizeof(T));
	}
};

template <class T>
class TypeToInt
{
	const static int v = -1;
};

template <> class TypeToInt<uchar>  { public: static const int v = GL_UNSIGNED_BYTE; };
template <> class TypeToInt<uint>   { public: static const int v = GL_UNSIGNED_INT; };
template <> class TypeToInt<ushort> { public: static const int v = GL_UNSIGNED_SHORT; };
template <> class TypeToInt<char>   { public: static const int v = GL_BYTE; };
template <> class TypeToInt<int>    { public: static const int v = GL_INT; };
template <> class TypeToInt<short>  { public: static const int v = GL_SHORT; };
template <> class TypeToInt<float>  { public: static const int v = GL_FLOAT; };

#endif